"""
构建向量检索索引
从数据库中的论文构建FAISS索引
"""
import logging
import sys
from pathlib import Path

from data_collection.data_storage import DataStorage
from retrieval import TextEmbedder, FAISSIndexer
from config import FAISS_INDEX_DIR

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def build_index():
    """构建FAISS索引"""
    logger.info("开始构建向量索引...")
    
    # 1. 加载数据
    storage = DataStorage()
    papers = storage.get_all_papers()
    
    if not papers:
        logger.error("数据库中没有论文数据，请先运行数据采集脚本")
        return
    
    logger.info(f"找到 {len(papers)} 篇论文")
    
    # 2. 向量化
    logger.info("开始向量化论文...")
    embedder = TextEmbedder()
    embeddings = embedder.encode_papers(papers)
    logger.info(f"向量化完成，形状: {embeddings.shape}")
    
    # 3. 构建索引
    logger.info("构建FAISS索引...")
    indexer = FAISSIndexer(dimension=embeddings.shape[1])
    pmids = [paper['pmid'] for paper in papers]
    indexer.add_vectors(embeddings, pmids)
    
    # 4. 保存索引
    index_path = FAISS_INDEX_DIR / "medical_index"
    FAISS_INDEX_DIR.mkdir(parents=True, exist_ok=True)
    indexer.save(str(index_path))
    
    logger.info("索引构建完成！")
    logger.info(f"索引保存到: {index_path}")


if __name__ == "__main__":
    build_index()
